//
// Created by zengao on 2022/12/29.
//

#include "Parser.h"
#include "AstNode.h"
#include <cassert>
#include "Diag.h"

using namespace C100;

std::shared_ptr<ProgramNode> C100::Parser::parse() {
    auto node = std::make_shared<ProgramNode>();
    locals = &node->localVars;
    while (lex.currentToken->kind != TokenKind::Eof) {
        node->stmts.push_back(parseStmt());
    }
    return node;
}

std::shared_ptr<AstNode> C100::Parser::parseExpr() {
    return parseAssignExpr();
}

std::shared_ptr<AstNode> C100::Parser::parseAddExpr() {
    std::shared_ptr<AstNode> left = parseMultiExpr();
    while (lex.currentToken->kind == TokenKind::Add || lex.currentToken->kind == TokenKind::Sub) {
        BinaryOperator anOperator = BinaryOperator::Add;
        if (lex.currentToken->kind == TokenKind::Sub) {
            anOperator = BinaryOperator::Sub;
        }
        lex.getNextToken();
        auto node = std::make_shared<BinaryNode>();
        node->binOp = anOperator;
        node->Lhs = left;
        node->Rhs = parseMultiExpr();
        left = node;
    }
    return left;
}

std::shared_ptr<AstNode> C100::Parser::parseMultiExpr() {
    std::shared_ptr<AstNode> left = parsePrimaryExpr();
    while (lex.currentToken->kind == TokenKind::Mul || lex.currentToken->kind == TokenKind::Div) {
        BinaryOperator anOperator = BinaryOperator::Mul;
        if (lex.currentToken->kind == TokenKind::Div) {
            anOperator = BinaryOperator::Div;
        }
        lex.getNextToken();
        auto node = std::make_shared<BinaryNode>();
        node->binOp = anOperator;
        node->Lhs = left;
        node->Rhs = parsePrimaryExpr();
        left = node;
    }
    return left;
}

std::shared_ptr<AstNode> C100::Parser::parsePrimaryExpr() {
    if (lex.currentToken->kind == TokenKind::LParent) {
        lex.getNextToken();
        auto node = parseExpr();
        lex.expectToken(TokenKind::RParent);
        return node;
    } else if (lex.currentToken->kind == TokenKind::Num) {
        auto node = std::make_shared<ConstantNode>();
        node->value = lex.currentToken->value;
        lex.getNextToken();
        return node;
    } else if (lex.currentToken->kind == TokenKind::Identifier) {
        auto node = std::make_shared<VarExprNode>();
        std::shared_ptr<Var> obj = findLocalVar(lex.currentToken->content);
        if (!obj) {
            obj = makeLocalVar(lex.currentToken->content);
        }
        node->varObj = obj;
        lex.getNextToken();
        return node;
    } else {
        diagE(lex.sourceCode, lex.currentToken->location.line, lex.currentToken->location.col, "not support node");
    }
    return nullptr;
}

std::shared_ptr<AstNode> Parser::parseStmt() {
    if (lex.currentToken->kind == TokenKind::If) {
        auto node = std::make_shared<IfStmtNode>();
        lex.getNextToken();
        lex.expectToken(TokenKind::LParent);
        node->Cond = parseExpr();
        lex.expectToken(TokenKind::RParent);
        node->Then = parseStmt();
        if (lex.currentToken->kind == TokenKind::Else) {
            lex.getNextToken();
            node->Else = parseStmt();
        }
        return node;
    } else if (lex.currentToken->kind == TokenKind::LBrace) {
        auto node = std::make_shared<BlockStmtNode>();
        lex.getNextToken();
        while (lex.currentToken->kind != TokenKind::RBrace){
            node->stmts.push_back(parseStmt());
        }
        lex.expectToken(TokenKind::RBrace);
        return node;
    } else {
        auto node = std::make_shared<ExprStmtNode>();
        node->Lhs = parseExpr();
        assert(lex.currentToken->kind == TokenKind::Semicolon);
        lex.getNextToken();
        return node;
    }
}

std::shared_ptr<AstNode> Parser::parseEqualExpr() {
    auto left = parseRelationalExpr();
    while (lex.currentToken->kind == TokenKind::Equal ||
           lex.currentToken->kind == TokenKind::PipeEqual) {
        BinaryOperator op = BinaryOperator::Equal;
        if (lex.currentToken->kind == TokenKind::PipeEqual) {
            op = BinaryOperator::PipeEqual;
        }
        lex.getNextToken();
        auto node = std::make_shared<BinaryNode>();
        node->binOp = op;
        node->Lhs = left;
        node->Rhs = parseRelationalExpr();
        left = node;
    }
    return left;
}

std::shared_ptr<AstNode> Parser::parseRelationalExpr() {
    auto left = parseAddExpr();
    while (lex.currentToken->kind == TokenKind::Greater ||
           lex.currentToken->kind == TokenKind::GreaterEqual ||
           lex.currentToken->kind == TokenKind::Lesser ||
           lex.currentToken->kind == TokenKind::LesserEqual) {
        BinaryOperator op = BinaryOperator::Greater;
        if (lex.currentToken->kind == TokenKind::GreaterEqual) {
            op = BinaryOperator::GreaterEqual;
        } else if (lex.currentToken->kind == TokenKind::Lesser) {
            op = BinaryOperator::Lesser;
        }
        if (lex.currentToken->kind == TokenKind::LesserEqual) {
            op = BinaryOperator::LesserEqual;
        }
        lex.getNextToken();
        auto node = std::make_shared<BinaryNode>();
        node->binOp = op;
        node->Lhs = left;
        node->Rhs = parseAddExpr();
        left = node;
    }
    return left;
}


std::shared_ptr<AstNode> Parser::parseAssignExpr() {
    auto left = parseEqualExpr();
    if (lex.currentToken->kind == TokenKind::Assign) {
        lex.getNextToken();
        auto node = std::make_shared<AssignExprNode>();
        node->Lhs = left;
        node->Rhs = parseAssignExpr();
        return node;
    }
    return left;
}

std::shared_ptr<Var> Parser::findLocalVar(std::string_view name) {
    if (localsMap.find(name) != localsMap.end()) {
        return localsMap[name];
    }
    return nullptr;
}

std::shared_ptr<Var> Parser::makeLocalVar(std::string_view name) {
    auto obj = std::make_shared<Var>();
    obj->name = name;
    obj->offset = 0;
    locals->push_back(obj);
    localsMap[name] = obj;
    return obj;
}

